package com.jboltai.capability.vdb;


import com.jboltai.capability.vdb.collection.VDBCollection;
import com.jboltai.capability.vdb.collection.field.VDBField;
import com.jboltai.capability.vdb.doc.VDBDoc;
import com.jboltai.capability.vdb.search.VDBConditions;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.resource.vdb.VDBResource;
import com.jboltai.resource.vdb.VDBResourceCenter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 向量数据库
 */
public abstract class JBoltVDB {

    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);
    static final String DEFAULT_DB_NAME = "main";

    /**
     * 获取腾讯向量数据库操作对象,通过默认名称：main获取
     * @return
     */
    public static JBoltVDB use() {
        return use(DEFAULT_DB_NAME);
    }


    private static ConcurrentHashMap<String, JBoltVDB> DBS = new ConcurrentHashMap<>();

    /**
     * 获取腾讯向量数据库操作对象
     * @param dbId
     * @return
     */
    public static JBoltVDB use(String dbId) {
        JBoltVDB vdb = DBS.get(dbId);
        if (vdb  == null) {
            VDBResource resource = VDBResourceCenter.get(dbId);
            if (resource == null ) {
                throw new RuntimeException("未注册id为[" + dbId + "]的向量数据库资源");
            }
            try {
                vdb = resource.getType().getvDbClass().getConstructor(VDBResource.class).newInstance(resource);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
            DBS.putIfAbsent(dbId, vdb);
        }
        return vdb;

    }

    /**
     * 删除数据库缓存
     * @param dbId
     */
    public static void remove(String dbId) {
        DBS.remove(dbId);
    }


    protected JBoltVDB(VDBResource resource) {
        this.resource = resource;
    }


    protected VDBResource resource;


    public VDBResource getResource() {
        return resource;
    }


    /**
     * 创建集合
     * @param name
     * @param fields 集合的字段，默认内置了一个name为id，type为string类型的主键字段，如果没有特殊需求，可以不用重复指定。
     *               如果fields中提供一个主键类型的字段，内置的id字段会自动取消。
     * @return
     */
    public VDBOptResult<Boolean> createCollection(String name, List<VDBField> fields) {
        if (fields == null || fields.isEmpty()) {
            throw new IllegalArgumentException("请指定集合的字段");
        }
        VDBCollection vDbCollection = new VDBCollection(this).setName(name);
        for (VDBField field : fields) {
            vDbCollection.addField(field);
        }
        return createCollection(vDbCollection);
    }
    /**
     * 创建集合
     * @param name
     */
    public VDBOptResult<Boolean> createCollection(String name, int shardNum, List<VDBField> fields) {
        if (fields == null || fields.isEmpty()) {
            throw new IllegalArgumentException("请指定集合的字段");
        }
        VDBCollection vDbCollection = new VDBCollection(this).setName(name);
        vDbCollection.setShardNum(shardNum);
        for (VDBField field : fields) {
            vDbCollection.addField(field);
        }
        return createCollection(vDbCollection);
    }

    /**
     * 创建集合
     * @param collection
     * @return
     */
    public abstract VDBOptResult<Boolean> createCollection(VDBCollection collection);

    /**
     * 删除集合
     * @param name
     * @return
     */
    public abstract VDBOptResult<Boolean> dropCollection(String name);

    /**
     * 检测某个集合是否已存在
     * @param name
     * @return
     */
    public abstract VDBOptResult<Boolean> checkCollectionExist(String name);


    /**
     * 清空集合内数据
     * @param name
     * @return
     */
    public abstract VDBOptResult<Boolean> truncateCollection(String name);


    /**
     * 插入数据
     * @param id
     * @param vectorText 要进行向量化的文本
     * @param fieldValues 字段值,按照 field、value、filed、value的顺序传入字段名和值
     * @return
     */
    public VDBOptResult<Integer> insert(String collectionName, String id, String vectorText, Object... fieldValues ) {
        if (fieldValues != null && fieldValues.length % 2 != 0) {
            throw new IllegalArgumentException("fieldValues参数必须成对传入。");
        }
        HashMap<String, Object> map = null;
        if (fieldValues != null) {
            map = new HashMap<>();
            for (int i = 0; i < fieldValues.length; i+=2) {
                map.put(fieldValues[i].toString(), fieldValues[i+1]);
            }
        }

        return insert(collectionName, id, vectorText, map);
    }

    /**
     * 插入数据
     * @param id
     * @param vectorText 要进行向量化的文本
     * @param fieldValues 字段值，key是字段名，value是字段值
     * @return
     */
    public  VDBOptResult<Integer> insert(String collectionName, String id, String vectorText, HashMap<String, Object> fieldValues) {

        return insert(collectionName, new VDBDoc(id, vectorText, fieldValues));
    }


    /**
     * 插入数据
     *
     * @param collectionName
     * @param doc
     * @return
     */
    public abstract VDBOptResult<Integer> insert(String collectionName, VDBDoc... doc);


    public VDBOptResult<Integer> insert(String collectionName, List<VDBDoc> docs) {
        if (docs == null || docs.isEmpty()) {
            throw new IllegalArgumentException("docs不能为空");
        }
        return insert(collectionName, docs.toArray(new VDBDoc[0]));
    }



    public VDBOptResult<VDBDoc> getById(String collectionName, String id) {
        VDBOptResult<List<VDBDoc>> result = search(collectionName, new VDBConditions().eq(VDBField.ID_FIELD_NAME, id));
        if (!result.isSuccess()) {
            return new VDBOptResult<>( result.getMsg(), result.getException());
        }
        return new VDBOptResult<>(true, (result.getResult() != null && !result.getResult().isEmpty()) ? result.getResult().get(0) : null);
    }

    /**
     * 查询数据
     */
    public abstract VDBOptResult<List<VDBDoc>> search(String collectionName, VDBConditions conditions);


    /**
     * 根据id删除数据
     * @param collectionName
     * @param id
     * @return
     */
    public VDBOptResult<Boolean> deleteById(String collectionName, String id) {
        return delete(collectionName, new VDBConditions().eq("id", id));
    }

    /**
     * 删除数据
     * @param collectionName
     * @param conditions
     * @return result代表影响的行数
     */
    public abstract VDBOptResult<Boolean> delete(String collectionName,  VDBConditions conditions);


    /**
     * 更新指定id的文本字段值
     * @param collectionName
     * @param id
     * @param doc
     * @return
     */
    public VDBOptResult<Integer> updateDocById(String collectionName, String id, String doc) {
        return update(collectionName, new VDBDoc().setDoc(doc), new VDBConditions().eq(VDBField.ID_FIELD_NAME, id));
    }

    /**
     * 更新指定id的文本字段和其他字段值
     * @param collectionName
     * @param id
     * @param doc
     * @param fieldValues
     * @return
     */
    public VDBOptResult<Integer> updateById(String collectionName, String id, String doc, Object... fieldValues) {
        if (fieldValues != null && fieldValues.length % 2 != 0) {
            throw new IllegalArgumentException("fieldValues参数必须成对传入。");
        }
        VDBConditions conditions = null;
        if (fieldValues != null && fieldValues.length > 0) {
            conditions = new VDBConditions();
            for (int i = 0; i < fieldValues.length; i+=2) {
                if (fieldValues[i+1] instanceof String) {
                    conditions.eq(fieldValues[i].toString(), (String)fieldValues[i+1]);
                } else if (fieldValues[i+1] instanceof Integer) {
                    conditions.eq(fieldValues[i].toString(), (Integer) fieldValues[i+1]);
                } else if (fieldValues[i+1] instanceof Long) {
                    conditions.eq(fieldValues[i].toString(), (Long)fieldValues[i+1]);
                }
                throw new RuntimeException("不支持的值类型：" + fieldValues[i+1]);
            }
        }

        return update(collectionName, new VDBDoc(id, doc), conditions);
    }

    /**
     * 根据doc 更新数据
     * @param collectionName
     * @param doc
     * @return
     */
    public VDBOptResult<Integer> update(String collectionName, VDBDoc doc) {
        return update(collectionName, doc, null);
    }

    /**
     * 更新文档
     * @param collectionName
     * @param doc 更新后的内容
     * @param conditions 更新的条件
     * @return result代表影响的行数
     */
    public abstract VDBOptResult<Integer> update(String collectionName, VDBDoc doc, VDBConditions conditions);



    /**
     * 获取数据库下的所有集合
     * @return
     */
    public abstract VDBOptResult<List<VDBCollection>> listCollections();

    /**
     * 创建向量数据库<br/>
     * <b>注意：目前该方法只在部分向量数据库中有效，比如腾讯向量数据库</b>
     * @param database
     */
    public abstract VDBOptResult<Void> createDatabase(String database);

    /**
     * 检测数据库是否存在</br>
     * <b>注意不是所有类型的vdb都支持该操作</b>
     * @param databse
     * @return
     */
    public abstract VDBOptResult<Boolean> checkDatabaseExist(String databse);

    /**
     * 删除数据库
     * @param database
     * @return
     */
    public abstract VDBOptResult<Boolean> dropDatabase(String database);
}
